# 导入必要的库
import streamlit as st
from langchain.vectorstores import Chroma
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
import os
from LLM import InternLM_LLM
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA
from ChatGLM_new import tongyi_embeddings,sspy_llm

def load_chain():
    # 加载问答链
    # 定义 Embeddings
    #embeddings = HuggingFaceEmbeddings(model_name="/root/data/model/sentence-transformer")
    embeddings = tongyi_embeddings

    # 向量数据库持久化路径
    persist_directory = 'data_base/vector_db/chroma'

    # 加载数据库
    vectordb = Chroma(
        persist_directory=persist_directory,  # 允许我们将persist_directory目录保存到磁盘上
        embedding_function=embeddings
    )

    #llm = InternLM_LLM(model_path="/root/data/model/Shanghai_AI_Laboratory/internlm-chat-7b")
    llm = sspy_llm
    # 你可以修改这里的 prompt template 来试试不同的问答效果
    template = """请使用以下提供的上下文来回答用户的问题。如果无法从上下文中得到答案，请回答你不知道，并总是使用中文回答。
    提供的上下文：
    ···
    {context}
    ···
    用户的问题: {question}
    你给的回答:"""

    QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"],
                                     template=template)

    # 运行 chain
    qa_chain = RetrievalQA.from_chain_type(llm,
                                           retriever=vectordb.as_retriever(),
                                           return_source_documents=True,
                                           chain_type_kwargs={"prompt": QA_CHAIN_PROMPT})

    return qa_chain

class ModelCenter:
    """
    存储问答 Chain 的对象 
    """
    def __init__(self):
        self.chain = load_chain()

    def qa_chain_self_answer(self, question: str, chat_history: list = []):
        """
        调用不带历史记录的问答链进行回答
        """
        if question == None or len(question) < 1:
            return "", chat_history
        try:
            result = self.chain({"query": question})["result"]
            chat_history.append(
                (question, result))
            return result, chat_history
        except Exception as e:
            return e, chat_history

model_center = ModelCenter()

st.title('InternLM')
st.subheader('书生浦语')

# 创建一个文本框组件，用于输入 prompt。
msg = st.text_input("请输入问题：")

# 创建一个聊天记录列表
chat_history = []

if 'messages' not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 当用户点击发送按钮或按下 Enter 键时
if msg:
    # 显示用户消息
    with st.chat_message("user"):
        st.markdown(msg)
    
    # 显示助理正在输入的状态
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        message_placeholder.markdown("思考中...")

    # 调用问答链并获取响应
    response, history = model_center.qa_chain_self_answer(msg, chat_history)
    st.session_state.messages.append({"role": "user", "content": msg})
    st.session_state.messages.append({"role": "assistant", "content": response})

    # 显示助理的回复
    with st.chat_message("assistant"):
        message_placeholder.markdown(response)



# 添加提醒信息
st.markdown("""
提醒：<br>
1. 初始化数据库时间可能较长，请耐心等待。
2. 使用中如果出现异常，将会在文本输入框进行展示，请不要惊慌。 <br>
""", unsafe_allow_html=True)

